#include <rclcpp/rclcpp.hpp>
#include <opentelemetry/trace/provider.h>
#include <opentelemetry/exporters/ostream/span_exporter.h>
#include <opentelemetry/sdk/trace/simple_processor.h>
#include <opentelemetry/sdk/trace/tracer_provider.h>
#include <opentelemetry/sdk/trace/samplers/always_on.h>
#include <opentelemetry/sdk/resource/resource.h>
#include <opentelemetry/nostd/shared_ptr.h>
#include <opentelemetry/nostd/unique_ptr.h>

#include <chrono>
#include <thread>
#include <string>
#include <vector>
#include <memory>
#include <random>

namespace trace_api = opentelemetry::trace;
namespace trace_sdk = opentelemetry::sdk::trace;
namespace resource_sdk = opentelemetry::sdk::resource;
namespace nostd = opentelemetry::nostd;

// Initialize OpenTelemetry tracing
void InitTracer() {
    auto exporter = std::unique_ptr<trace_sdk::SpanExporter>(
        new opentelemetry::exporter::trace::OStreamSpanExporter);
    
    auto processor = std::unique_ptr<trace_sdk::SpanProcessor>(
        new trace_sdk::SimpleSpanProcessor(std::move(exporter)));
    
    auto resource = resource_sdk::Resource::Create({
        {"service.name", "ros2-function-tracing-demo"},
        {"service.version", "1.0.0"},
        {"service.instance.id", "function-tracing-001"}
    });
    
    auto sampler = std::unique_ptr<trace_sdk::Sampler>(
        new trace_sdk::AlwaysOnSampler);
    
    auto provider = nostd::shared_ptr<trace_api::TracerProvider>(
        new trace_sdk::TracerProvider(std::move(processor), resource, std::move(sampler)));
    
    trace_api::Provider::SetTracerProvider(provider);
}

class FunctionTracingDemo : public rclcpp::Node {
public:
    FunctionTracingDemo() : Node("function_tracing_demo") {
        // Initialize OpenTelemetry
        InitTracer();
        tracer_ = trace_api::Provider::GetTracerProvider()->GetTracer("function-tracing-demo");
        
        RCLCPP_INFO(this->get_logger(), "Function Tracing Demo started");
        RCLCPP_INFO(this->get_logger(), "Demonstrating single-process function call tracing...");
        
        // Run different function call scenarios
        //runBasicFunctionCallDemo();
        //runNestedFunctionCallDemo();
        runConcurrentFunctionCallDemo();
        //runErrorHandlingDemo();
        //runPerformanceAnalysisDemo();
    }

private:
    // Basic function call tracing
    void runBasicFunctionCallDemo() {
        RCLCPP_INFO(this->get_logger(), "\n=== Basic Function Call Demo ===");
        
        auto span = tracer_->StartSpan("basic_function_demo", {
            {"demo.type", "basic_function_calls"},
            {"demo.description", "Simple function call chain"}
        });
        
        auto scope = tracer_->WithActiveSpan(span);
        
        try {
            // Call a simple function
            int result = processUserRequest("user123", "get_profile");
            span->SetAttribute("function.result", result);
            span->AddEvent("basic_demo_completed");
            span->SetStatus(trace_api::StatusCode::kOk, "Basic function demo completed successfully");
            
            RCLCPP_INFO(this->get_logger(), "Basic function demo result: %d", result);
            
        } catch (const std::exception& e) {
            span->SetStatus(trace_api::StatusCode::kError, e.what());
            span->SetAttribute("error.message", e.what());
            RCLCPP_ERROR(this->get_logger(), "Error in basic function demo: %s", e.what());
        }
        
        span->End();
    }
    
    // Nested function call tracing
    void runNestedFunctionCallDemo() {
        RCLCPP_INFO(this->get_logger(), "\n=== Nested Function Call Demo ===");
        
        auto span = tracer_->StartSpan("nested_function_demo", {
            {"demo.type", "nested_function_calls"},
            {"demo.description", "Deep nested function call chain"}
        });
        
        auto scope = tracer_->WithActiveSpan(span);
        
        try {
            // Call a function that internally calls other functions
            std::string result = processComplexRequest("complex_request_001");
            span->SetAttribute("function.result", result);
            span->AddEvent("nested_demo_completed");
            span->SetStatus(trace_api::StatusCode::kOk, "Nested function demo completed successfully");
            
            RCLCPP_INFO(this->get_logger(), "Nested function demo result: %s", result.c_str());
            
        } catch (const std::exception& e) {
            span->SetStatus(trace_api::StatusCode::kError, e.what());
            span->SetAttribute("error.message", e.what());
            RCLCPP_ERROR(this->get_logger(), "Error in nested function demo: %s", e.what());
        }
        
        span->End();
    }
    
    // Concurrent function call tracing
    void runConcurrentFunctionCallDemo() {
        RCLCPP_INFO(this->get_logger(), "\n=== Concurrent Function Call Demo ===");
        
        auto span = tracer_->StartSpan("concurrent_function_demo", {
            {"demo.type", "concurrent_function_calls"},
            {"demo.description", "Multiple concurrent function calls"}
        });
        
        auto scope = tracer_->WithActiveSpan(span);
        
        try {
            std::vector<std::thread> threads;
            std::vector<int> results(3);
            
            // Launch concurrent function calls
            for (int i = 0; i < 3; ++i) {
                threads.emplace_back([this, i, &results]() {
                    results[i] = processConcurrentRequest("concurrent_request_" + std::to_string(i));
                });
            }
            
            // Wait for all threads to complete
            for (auto& thread : threads) {
                thread.join();
            }
            
            // Log results
            for (int i = 0; i < 3; ++i) {
                RCLCPP_INFO(this->get_logger(), "Concurrent request %d result: %d", i, results[i]);
            }
            
            span->SetAttribute("concurrent.results_count", static_cast<int64_t>(results.size()));
            span->AddEvent("concurrent_demo_completed");
            span->SetStatus(trace_api::StatusCode::kOk, "Concurrent function demo completed successfully");
            
        } catch (const std::exception& e) {
            span->SetStatus(trace_api::StatusCode::kError, e.what());
            span->SetAttribute("error.message", e.what());
            RCLCPP_ERROR(this->get_logger(), "Error in concurrent function demo: %s", e.what());
        }
        
        span->End();
    }
    
    // Error handling demo
    void runErrorHandlingDemo() {
        RCLCPP_INFO(this->get_logger(), "\n=== Error Handling Demo ===");
        
        auto span = tracer_->StartSpan("error_handling_demo", {
            {"demo.type", "error_handling"},
            {"demo.description", "Function call with error handling"}
        });
        
        auto scope = tracer_->WithActiveSpan(span);
        
        try {
            // Call a function that might throw an error
            int result = processRequestWithError("error_test_request");
            span->SetAttribute("function.result", result);
            span->AddEvent("error_handling_demo_completed");
            span->SetStatus(trace_api::StatusCode::kOk, "Error handling demo completed successfully");
            
            RCLCPP_INFO(this->get_logger(), "Error handling demo result: %d", result);
            
        } catch (const std::exception& e) {
            span->SetStatus(trace_api::StatusCode::kError, e.what());
            span->SetAttribute("error.message", e.what());
            span->SetAttribute("error.type", "std::exception");
            span->AddEvent("error_occurred", {
                {"error.details", e.what()}
            });
            RCLCPP_ERROR(this->get_logger(), "Error in error handling demo: %s", e.what());
        }
        
        span->End();
    }
    
    // Performance analysis demo
    void runPerformanceAnalysisDemo() {
        RCLCPP_INFO(this->get_logger(), "\n=== Performance Analysis Demo ===");
        
        auto span = tracer_->StartSpan("performance_analysis_demo", {
            {"demo.type", "performance_analysis"},
            {"demo.description", "Function call performance analysis"}
        });
        
        auto scope = tracer_->WithActiveSpan(span);
        
        try {
            // Measure performance of different function calls
            auto start_time = std::chrono::high_resolution_clock::now();
            
            int result1 = fastFunction("fast_request");
            int result2 = slowFunction("slow_request");
            int result3 = mediumFunction("medium_request");
            
            auto end_time = std::chrono::high_resolution_clock::now();
            auto duration = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);
            
            span->SetAttribute("performance.total_duration_us", static_cast<int64_t>(duration.count()));
            span->SetAttribute("performance.fast_result", result1);
            span->SetAttribute("performance.slow_result", result2);
            span->SetAttribute("performance.medium_result", result3);
            span->AddEvent("performance_analysis_completed");
            span->SetStatus(trace_api::StatusCode::kOk, "Performance analysis demo completed successfully");
            
            RCLCPP_INFO(this->get_logger(), "Performance analysis results:");
            RCLCPP_INFO(this->get_logger(), "  Fast function result: %d", result1);
            RCLCPP_INFO(this->get_logger(), "  Slow function result: %d", result2);
            RCLCPP_INFO(this->get_logger(), "  Medium function result: %d", result3);
            RCLCPP_INFO(this->get_logger(), "  Total duration: %ld μs", duration.count());
            
        } catch (const std::exception& e) {
            span->SetStatus(trace_api::StatusCode::kError, e.what());
            span->SetAttribute("error.message", e.what());
            RCLCPP_ERROR(this->get_logger(), "Error in performance analysis demo: %s", e.what());
        }
        
        span->End();
    }
    
    // Function implementations with tracing
    
    int processUserRequest(const std::string& user_id, const std::string& action) {
        auto span = tracer_->StartSpan("process_user_request", {
            {"function.name", "processUserRequest"},
            {"user.id", user_id},
            {"action.type", action}
        });
        
        auto scope = tracer_->WithActiveSpan(span);
        
        try {
            span->AddEvent("validating_user_request");
            
            // Simulate validation
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
            
            if (user_id.empty() || action.empty()) {
                span->SetStatus(trace_api::StatusCode::kError, "Invalid request parameters");
                span->SetAttribute("error.code", "INVALID_PARAMS");
                span->End();
                return -1;
            }
            
            span->AddEvent("request_validated");
            
            // Simulate processing
            std::this_thread::sleep_for(std::chrono::milliseconds(20));
            
            int result = validateUser(user_id) + executeAction(action);
            
            span->SetAttribute("function.result", result);
            span->AddEvent("request_processed");
            span->SetStatus(trace_api::StatusCode::kOk, "User request processed successfully");
            span->End();
            
            return result;
            
        } catch (const std::exception& e) {
            span->SetStatus(trace_api::StatusCode::kError, e.what());
            span->SetAttribute("error.message", e.what());
            span->End();
            throw;
        }
    }
    
    std::string processComplexRequest(const std::string& request_id) {
        auto span = tracer_->StartSpan("func1_process_request", {
            {"function.name", "processComplexRequest"},
            {"function.level", "1"},
            {"request.id", request_id}
        });
        
        auto scope = tracer_->WithActiveSpan(span);
        
        try {
            span->AddEvent("func1_started");
            
            // Call func2
            std::this_thread::sleep_for(std::chrono::milliseconds(20));
            std::string result = func2_validateRequest(request_id);
            
            span->SetAttribute("function.result", result);
            span->AddEvent("func1_completed");
            span->SetStatus(trace_api::StatusCode::kOk, "Func1 completed successfully");
            span->End();
            
            return result;
            
        } catch (const std::exception& e) {
            span->SetStatus(trace_api::StatusCode::kError, e.what());
            span->SetAttribute("error.message", e.what());
            span->End();
            throw;
        }
    }
    
    int processConcurrentRequest(const std::string& request_id) {
        auto span = tracer_->StartSpan("process_concurrent_request", {
            {"function.name", "processConcurrentRequest"},
            {"request.id", request_id}
        });
        
        auto scope = tracer_->WithActiveSpan(span);
        
        try {
            span->AddEvent("starting_concurrent_processing");
            
            // Simulate some work
            std::this_thread::sleep_for(std::chrono::milliseconds(50));
            
            // Generate a result based on request ID
            int result = std::hash<std::string>{}(request_id) % 1000;
            
            span->SetAttribute("function.result", result);
            span->AddEvent("concurrent_processing_completed");
            span->SetStatus(trace_api::StatusCode::kOk, "Concurrent request processed successfully");
            span->End();
            
            return result;
            
        } catch (const std::exception& e) {
            span->SetStatus(trace_api::StatusCode::kError, e.what());
            span->SetAttribute("error.message", e.what());
            span->End();
            throw;
        }
    }
    
    int processRequestWithError(const std::string& request_id) {
        auto span = tracer_->StartSpan("process_request_with_error", {
            {"function.name", "processRequestWithError"},
            {"request.id", request_id}
        });
        
        auto scope = tracer_->WithActiveSpan(span);
        
        try {
            span->AddEvent("starting_error_test_processing");
            
            // Simulate some processing
            std::this_thread::sleep_for(std::chrono::milliseconds(20));
            
            // Simulate an error condition
            if (request_id.find("error") != std::string::npos) {
                span->SetStatus(trace_api::StatusCode::kError, "Simulated error for testing");
                span->SetAttribute("error.code", "SIMULATED_ERROR");
                span->AddEvent("error_simulated");
                span->End();
                throw std::runtime_error("Simulated error for testing purposes");
            }
            
            int result = 42; // Success case
            
            span->SetAttribute("function.result", result);
            span->AddEvent("error_test_processing_completed");
            span->SetStatus(trace_api::StatusCode::kOk, "Request processed successfully");
            span->End();
            
            return result;
            
        } catch (const std::exception& e) {
            span->SetStatus(trace_api::StatusCode::kError, e.what());
            span->SetAttribute("error.message", e.what());
            span->End();
            throw;
        }
    }
    
    int fastFunction(const std::string& request) {
        auto span = tracer_->StartSpan("fast_function", {
            {"function.name", "fastFunction"},
            {"function.type", "fast"},
            {"request", request}
        });
        
        auto scope = tracer_->WithActiveSpan(span);
        
        try {
            span->AddEvent("fast_processing_started");
            
            // Simulate fast processing
            std::this_thread::sleep_for(std::chrono::milliseconds(5));
            
            int result = 100;
            
            span->SetAttribute("function.result", result);
            span->SetAttribute("function.duration_ms", 5);
            span->AddEvent("fast_processing_completed");
            span->SetStatus(trace_api::StatusCode::kOk, "Fast function completed");
            span->End();
            
            return result;
            
        } catch (const std::exception& e) {
            span->SetStatus(trace_api::StatusCode::kError, e.what());
            span->SetAttribute("error.message", e.what());
            span->End();
            throw;
        }
    }
    
    int slowFunction(const std::string& request) {
        auto span = tracer_->StartSpan("slow_function", {
            {"function.name", "slowFunction"},
            {"function.type", "slow"},
            {"request", request}
        });
        
        auto scope = tracer_->WithActiveSpan(span);
        
        try {
            span->AddEvent("slow_processing_started");
            
            // Simulate slow processing
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
            
            int result = 200;
            
            span->SetAttribute("function.result", result);
            span->SetAttribute("function.duration_ms", 100);
            span->AddEvent("slow_processing_completed");
            span->SetStatus(trace_api::StatusCode::kOk, "Slow function completed");
            span->End();
            
            return result;
            
        } catch (const std::exception& e) {
            span->SetStatus(trace_api::StatusCode::kError, e.what());
            span->SetAttribute("error.message", e.what());
            span->End();
            throw;
        }
    }
    
    int mediumFunction(const std::string& request) {
        auto span = tracer_->StartSpan("medium_function", {
            {"function.name", "mediumFunction"},
            {"function.type", "medium"},
            {"request", request}
        });
        
        auto scope = tracer_->WithActiveSpan(span);
        
        try {
            span->AddEvent("medium_processing_started");
            
            // Simulate medium processing
            std::this_thread::sleep_for(std::chrono::milliseconds(30));
            
            int result = 150;
            
            span->SetAttribute("function.result", result);
            span->SetAttribute("function.duration_ms", 30);
            span->AddEvent("medium_processing_completed");
            span->SetStatus(trace_api::StatusCode::kOk, "Medium function completed");
            span->End();
            
            return result;
            
        } catch (const std::exception& e) {
            span->SetStatus(trace_api::StatusCode::kError, e.what());
            span->SetAttribute("error.message", e.what());
            span->End();
            throw;
        }
    }
    
    // Simple nested function calls: func1 -> func2 -> func3
    std::string func2_validateRequest(const std::string& request_id) {
        auto span = tracer_->StartSpan("func2_validate_request", {
            {"function.name", "func2_validateRequest"},
            {"function.level", "2"},
            {"request.id", request_id}
        });
        
        auto scope = tracer_->WithActiveSpan(span);
        
        try {
            span->AddEvent("func2_started");
            
            // Call func3
            std::this_thread::sleep_for(std::chrono::milliseconds(15));
            std::string result = func3_finalizeRequest(request_id);
            
            span->SetAttribute("function.result", result);
            span->AddEvent("func2_completed");
            span->SetStatus(trace_api::StatusCode::kOk, "Func2 completed successfully");
            span->End();
            
            return result;
            
        } catch (const std::exception& e) {
            span->SetStatus(trace_api::StatusCode::kError, e.what());
            span->SetAttribute("error.message", e.what());
            span->End();
            throw;
        }
    }
    
    std::string func3_finalizeRequest(const std::string& request_id) {
        auto span = tracer_->StartSpan("func3_finalize_request", {
            {"function.name", "func3_finalizeRequest"},
            {"function.level", "3"},
            {"request.id", request_id}
        });
        
        auto scope = tracer_->WithActiveSpan(span);
        
        try {
            span->AddEvent("func3_started");
            
            // Final processing
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
            std::string result = "finalized_" + request_id;
            
            span->SetAttribute("function.result", result);
            span->AddEvent("func3_completed");
            span->SetStatus(trace_api::StatusCode::kOk, "Func3 completed successfully");
            span->End();
            
            return result;
            
        } catch (const std::exception& e) {
            span->SetStatus(trace_api::StatusCode::kError, e.what());
            span->SetAttribute("error.message", e.what());
            span->End();
            throw;
        }
    }

    // Helper functions
    int validateUser(const std::string& user_id) {
        auto span = tracer_->StartSpan("validate_user", {
            {"function.name", "validateUser"},
            {"user.id", user_id}
        });
        
        auto scope = tracer_->WithActiveSpan(span);
        
        try {
            span->AddEvent("user_validation_started");
            
            // Simulate user validation
            std::this_thread::sleep_for(std::chrono::milliseconds(8));
            
            int result = user_id.length() * 10; // Simple validation logic
            
            span->SetAttribute("validation.result", result);
            span->AddEvent("user_validation_completed");
            span->SetStatus(trace_api::StatusCode::kOk, "User validation completed");
            span->End();
            
            return result;
            
        } catch (const std::exception& e) {
            span->SetStatus(trace_api::StatusCode::kError, e.what());
            span->SetAttribute("error.message", e.what());
            span->End();
            throw;
        }
    }
    
    int executeAction(const std::string& action) {
        auto span = tracer_->StartSpan("execute_action", {
            {"function.name", "executeAction"},
            {"action.type", action}
        });
        
        auto scope = tracer_->WithActiveSpan(span);
        
        try {
            span->AddEvent("action_execution_started");
            
            // Simulate action execution
            std::this_thread::sleep_for(std::chrono::milliseconds(12));
            
            int result = action.length() * 5; // Simple action logic
            
            span->SetAttribute("action.result", result);
            span->AddEvent("action_execution_completed");
            span->SetStatus(trace_api::StatusCode::kOk, "Action execution completed");
            span->End();
            
            return result;
            
        } catch (const std::exception& e) {
            span->SetStatus(trace_api::StatusCode::kError, e.what());
            span->SetAttribute("error.message", e.what());
            span->End();
            throw;
        }
    }
    

    nostd::shared_ptr<trace_api::Tracer> tracer_;
};

int main(int argc, char* argv[]) {
    rclcpp::init(argc, argv);
    
    auto node = std::make_shared<FunctionTracingDemo>();
    
    RCLCPP_INFO(node->get_logger(), "Starting Function Tracing Demo...");
    RCLCPP_INFO(node->get_logger(), "This demo shows single-process function call tracing with OpenTelemetry");
    
    // Keep the node alive to see all demo results
    std::this_thread::sleep_for(std::chrono::seconds(3));
    
    RCLCPP_INFO(node->get_logger(), "\n=== Function Tracing Demo Complete ===");
    RCLCPP_INFO(node->get_logger(), "All function call tracing demos completed. Check the trace output above.");
    
    rclcpp::shutdown();
    
    return 0;
}
